﻿using System;
using System.Collections.Generic;
using ASquared;
using ASquared.SymbolicMath;
using ASquared.ModelOptimization;
using Csu.Modsim.ModsimModel;
using Microsoft.SolverFoundation.Common;
using Microsoft.SolverFoundation.Services;

namespace iwpm.optimization
{
    public enum PowerSystemSolverMethod { ConditionalGradient, Simplex, Default }

    public class PowerSystemBase : IConditionalGradientSolvableFunction
    {
        public enum ObjectiveType { Lagrangian, AugmentedLagrangian };
        public enum ScaleType { SqrDiff, SqrDiffOverCap, InvNSCE };

        // Instance variables
        private Csu.Modsim.ModsimModel.Model _model;
        private double[][,] _origTargets;
        private HydropowerTarget[] _targets;
        private Symbol _fxn;
        private Matrix _c, _A, _b, _origCosts;
        private Matrix _decisions;
        private Matrix _mask;
        private List<VariableType> _vTypeList;
        private List<ConstraintType> _cTypeList;
        private String[] _vars;
        private Symbol[] _symVars;
        private bool _isConverged;
        private double _weight;
        private ObjectiveType _type = ObjectiveType.AugmentedLagrangian;
        private PowerSystemSolverMethod _solverMethod = PowerSystemSolverMethod.Default;
        private LinearSolver _linearSolver;
        private LagrangeMultiplierMatrix _mu;
        private SolverContext _context = SolverContext.GetContext();
        private Microsoft.SolverFoundation.Services.Model _solverModel = null;
        private double _load;
        private List<double> _caps;
        public ScaleType scaleoption = ScaleType.SqrDiffOverCap;
        private double[] _meantargets; 

        // events
        public event EventHandler Initializing; 
        public event EventHandler Initialized;
        private EventArgs _emptyArgs = new EventArgs();

        // Properties
        public double[] meantargets { get { return _meantargets; } } 
        public Matrix variableMask { get { return _mask; } }
        public Matrix costs { get { return _c; } set { _c = value; } }
        public Matrix decisions { get { return _decisions; } set { _decisions = value; } }
        public Matrix A { get { return _A; } set { _A = value; } }
        public Matrix b { get { return _b; } set { _b = value; } }
        public VariableType[] VariableTypes { get { return _vTypeList.ToArray(); } set { _vTypeList = new List<VariableType>(value); } }
        public List<VariableType> VariableTypesList { get { return _vTypeList; } set { _vTypeList = value; } }
        public ConstraintType[] ConstraintTypes { get { return _cTypeList.ToArray(); } set { _cTypeList = new List<ConstraintType>(value); } }
        public List<ConstraintType> ConstraintTypesList { get { return _cTypeList; } set { _cTypeList = value; } }
        public List<double> Capacities { get { return _caps; } set { _caps = value; } }
        public String[] variables { get { return _vars; } }
        public bool IsConverged { get { return _isConverged; } set { _isConverged = value; } }
        public double[][,] OriginalTargets { get { return _origTargets; } set { _origTargets = value; } }
        public double Weight { get { return _weight; } set { _weight = value; } }
        public double Load { get { return _load; } set { _load = value; } }
        public bool MinimizeSwing { get; set; }
        public double TimestepHours { get { return _model.timeStep.ToTimeSpan(_model.mInfo.CurrentBegOfPeriodDate).TotalHours; } }
        public PowerSystemSolverMethod SolverMethod { get { return _solverMethod; } set { _solverMethod = value; } }
        public LinearSolver LinearSolver { get { return _linearSolver; } set { _linearSolver = value; } }
        public LagrangeMultiplierMatrix Mu { get { return _mu; } }
        public ObjectiveType ObjType { get { return _type; } set { _type = value; } }
        public Microsoft.SolverFoundation.Services.Model SolverModel { get { return _solverModel; } set { _solverModel = value; } }
        public double[] CurrentOriginalTargets
        {
            get
            {
                double[] targs = new double[_origTargets.Length];
                for (int i = 0; i < _origTargets.Length; i++)
                    targs[i] = _origTargets[i][_model.mInfo.CurrentModelTimeStepIndex, 0];
                return targs;
            }
        }

        public Symbol fxn
        {
            get
            {
                _fxn = 0.0;

                for (int i = 0; i < _targets.Length; i++)
                {
                    Symbol diff = (_symVars[i] - getOrigTarget(i));
                    // attempt to meet original hydropower targets
                    _fxn += _weight * (diff ^ 2);

                    // add the Lagrange portion 
                    if (_type == ObjectiveType.AugmentedLagrangian)
                    {
                        diff = (_targets[i].Energy - _symVars[i]);
                        _fxn += _mu[i] * diff;
                        _fxn += _mu.Alpha / 2.0 * (diff ^ 2);
                    }
                    else if (_type == ObjectiveType.Lagrangian)
                    {
                        _fxn += -_mu[i] * _symVars[i];
                    }
                }

                return _fxn;
            }
        }

        // Constructors 
        public PowerSystemBase(Csu.Modsim.ModsimModel.Model model)
            : this(model, null)
        {
        }

        public PowerSystemBase(Csu.Modsim.ModsimModel.Model model, HydropowerTarget[] targets)
        {
            _model = model;
            if (targets == null)
                _targets = _model.hydro.HydroTargets;
            else
                _targets = targets;

            // fill the original targets
            _origTargets = new double[_targets.Length][,];
            for (int i = 0; i < _targets.Length; i++)
            {
                double[,] t = _targets[i].EnergyTargets;
                int nrows = t.GetLength(0);
                int ncols = t.GetLength(1);
                _origTargets[i] = new double[nrows, ncols];
                for (int c = 0; c < ncols; c++)
                    for (int r = 0; r < nrows; r++)
                        _origTargets[i][r, c] = t[r, c];
            }
        }

        // Methods 
        public void Initialize()
        {
            if (this.Initializing != null)
                this.Initializing(this, _emptyArgs);

            // initialize instance variables
            _isConverged = false;
            _vars = new string[_targets.Length];
            _symVars = new Symbol[_targets.Length];
            _mask = new Matrix(_targets.Length, (double)1.0);
            _vTypeList = new List<VariableType>(_targets.Length);
            _cTypeList = new List<ConstraintType>(_targets.Length);
            _c = new Matrix(_targets.Length, 1, 1.0);
            
            // get the means
            _meantargets = new double[_targets.Length]; 
            for (int i = 0; i< _targets.Length; i++) 
            {
                double sum = 0.0; 
                for (int t = 0; t < _origTargets[i].GetLength(0); t++) 
                    sum += _origTargets[i][t, 0]; 
                _meantargets[i] = sum/_origTargets[i].GetLength(0);
            }

            // Add the first constraint (generation = load)
            _cTypeList.Add(ConstraintType.Equality);
            _A = new Matrix(1, _targets.Length, 1.0);
            _b = new Matrix(1, 1, HydroTargetUpdater.SumEnergy(_targets, true, false));

            // build linear model for solver 
            _context.ClearModel();
            _solverModel = _context.CreateModel();
            List<Decision> Ps = new List<Decision>();
            _caps = new List<double>(); 
            Decision P;
            Term goal = 0.0, top = 0.0, bottom = 0.0;

            // Add the rest of the constraints and targets
            for (int i = 0; i < _targets.Length; i++)
            {
                _vars[i] = "P_" + i;
                _symVars[i] = _vars[i];

                // constraints and variable types 
                if (_targets[i].HasOnlyPumps)
                {
                    _caps.Add(_targets[i].PowerCapacityMW * TimestepHours); 
                    P = new Decision(Domain.RealRange(-_targets[i].PowerCapacityMW * TimestepHours, 0.0), _vars[i]);

                    this.AddConstraint(i, 1.0, -_targets[i].PowerCapacityMW * TimestepHours, ConstraintType.LeftGreaterThanRight);
                    _vTypeList.Add(VariableType.Negative);
                }
                else if (_targets[i].HasOnlyTurbines)
                {
                    _caps.Add(_targets[i].PowerCapacityMW * TimestepHours);
                    P = new Decision(Domain.RealRange(0, _targets[i].PowerCapacityMW * TimestepHours), _vars[i]);

                    this.AddConstraint(i, 1.0, _targets[i].PowerCapacityMW * TimestepHours, ConstraintType.LeftLessThanRight);
                    _vTypeList.Add(VariableType.NonNegative);
                }
                else
                {
                    _caps.Add(Math.Max(_targets[i].DownReserveCapacityMW, _targets[i].UpReserveCapacityMW) * TimestepHours); 
                    P = new Decision(Domain.RealRange(-_targets[i].DownReserveCapacityMW * TimestepHours, _targets[i].UpReserveCapacityMW * TimestepHours), _vars[i]);

                    this.AddConstraint(i, 1.0, _targets[i].UpReserveCapacityMW * TimestepHours, ConstraintType.LeftLessThanRight);
                    this.AddConstraint(i, 1.0, -_targets[i].DownReserveCapacityMW * TimestepHours, ConstraintType.LeftGreaterThanRight);
                    _vTypeList.Add(VariableType.Free);
                }

                _solverModel.AddDecision(P);
                Ps.Add(P);
            }

            // Set up Lagrange multipliers
            if (_mu == null || _mu.Length != _targets.Length)
            {
                _mu = new LagrangeMultiplierMatrix(_targets.Length);
                _mu.Alpha = _weight; 
                _mu.UseAugmented = (_type == ObjectiveType.AugmentedLagrangian);
            }

            // Set up SolverFoundation problem if desired 
            switch (_solverMethod)
            {
                case PowerSystemSolverMethod.Simplex:
                    _c = new Matrix(_targets.Length, (double)0.0);
                    for (int i = 0; i < _targets.Length; i++)
                    {
                        Term excess = Ps[i];
                        P = new Decision(Domain.Real, "s_plus_" + i);
                        _solverModel.AddDecision(P);
                        goal += _weight * P;
                        excess -= P;
                        P = new Decision(Domain.RealNonnegative, "s_neg_" + i);
                        _solverModel.AddDecision(P);
                        goal += _weight * P;
                        excess += P;
                        _solverModel.AddConstraint("excess_" + i, excess == getOrigTarget(i));

                        // s_up -> current target needs to go up
                        int s_up = this.AddVariable(_weight, VariableType.NonNegative);
                        // s_down -> current target needs to go down
                        int s_down = this.AddVariable(_weight, VariableType.NonNegative);
                        int[] indices = new int[] { i, s_up, s_down };
                        double[] weights = new double[] { 1.0, 1.0, -1.0 };
                        this.AddConstraint(indices, weights, getOrigTarget(i), ConstraintType.Equality);
                    }

                    // swing import 
                    Decision sImport = new Decision(Domain.RealNonnegative, "swingImport");
                    _solverModel.AddDecision(sImport);
                    Ps.Add(sImport);
                    goal += _weight * 100.0 * sImport;

                    // swing export 
                    Decision sExport = new Decision(Domain.RealNonnegative, "swingExport");
                    _solverModel.AddDecision(sExport);
                    goal += _weight * 100.0 * sExport;

                    // Pgen = Pload 
                    Term total = 0.0;
                    for (int i = 0; i < Ps.Count; i++)
                        total += Ps[i];
                    total -= sExport;
                    _solverModel.AddConstraint("EnergyBalance", total == _load);

                    // set objective function
                    _solverModel.AddGoal("EnergyTarget", GoalKind.Minimize, goal);

                    //// Add swing generator
                    //int swing_export = this.AddVariable(_weight * 100.0, VariableType.NonNegative);
                    //int swing_import = this.AddVariable(_weight * 100.0, VariableType.NonNegative); 
                    //_A[0, swing_export] = -1.0;
                    //_A[0, swing_import] = 1.0; 

                    break;
                    
                case PowerSystemSolverMethod.Default:
                    // add the Lagrange portion 
                    // Add bus constraints (Power In = Power Out) 
                    double mincap = double.MaxValue;
                    for (int i = 0; i < _targets.Length; i++)
                    {
                        // The main objective... to minimize the difference between the 
                        //   original targets and current targets, but still satisfy energy balance
                        double currcap = _caps[i]; 
                        switch (scaleoption)
                        {
                            case ScaleType.SqrDiff:
                                currcap = 1;
                                break;
                            case ScaleType.SqrDiffOverCap:
                                break;
                            case ScaleType.InvNSCE:
                                currcap = 1;
                                bottom += Microsoft.SolverFoundation.Services
                                    .Model.Power((getOrigTarget(i) - _meantargets[i]),2);
                                break;
                        }

                        if (mincap > currcap) mincap = currcap;
                        top += Microsoft.SolverFoundation.Services
                            .Model.Power((Ps[i] - getOrigTarget(i)) / currcap, 2);

                        // Don't want excess variable to be deviating from zero
                        if (_type == ObjectiveType.AugmentedLagrangian)
                        {
                            Term diff = (_targets[i].Energy - Ps[i]); // / _caps[i];
                            goal += _mu[i] * diff; 
                            goal += _mu.Alpha / 2.0 * Microsoft.SolverFoundation.Services.Model.Power(diff, 2); 
                        }
                        else if (_type == ObjectiveType.Lagrangian)
                        {
                            goal -= _mu[i] * Ps[i]; // / _caps[i];
                        }
                    }
                    switch (scaleoption) 
                    {
                        case ScaleType.SqrDiff: 
                            goal += _weight * top;
                            break;
                        case ScaleType.SqrDiffOverCap:
                            goal += _weight * top;
                            break;
                        case ScaleType.InvNSCE:
                            goal += _weight * top / bottom;
                            break;
                    }

                    // Add swing bus variable
                    Decision s_r = new Decision(Domain.Real, "slack_right");
                    Decision s_l = new Decision(Domain.Real, "slack_left");
                    _solverModel.AddDecision(s_r);
                    _solverModel.AddDecision(s_l);
                    Ps.Add(s_r);
                    Ps.Add(s_l);

                    // Pgen = Pload 
                    Term Pgen = 0.0;
                    for (int i = 0; i < Ps.Count; i++)
                        Pgen += Ps[i];
                    _solverModel.AddConstraint("EnergyBalance", Pgen == _load);

                    // Set objective function... don't want imports or exports or to be deviating from original power targets
                    if (this.MinimizeSwing)
                        goal += _weight * (Microsoft.SolverFoundation.Services.Model.Power(s_l / mincap, 2) 
                                           + Microsoft.SolverFoundation.Services.Model.Power(s_r / mincap, 2));
                        // goal += _weight * 10.0 * Microsoft.SolverFoundation.Services.Model.Power(swing / mincap, 2);
                    _solverModel.AddGoal("EnergyTarget", GoalKind.Minimize, goal);

                    break;

                default:

                    break;
            }

            // set the original decisions and costs.
            _decisions = this.EnergyTargets();
            _origCosts = _c.Copy();

            if (this.Initialized != null)
                this.Initialized(this, _emptyArgs);
        }

        public double getOrigTarget(int targetIndex)
        {
            return _origTargets[targetIndex][_model.mInfo.CurrentModelTimeStepIndex, 0];
        }

        public double getOrigTarget(int[] targetIndices)
        {
            double sum = 0.0;
            foreach (int i in targetIndices)
                sum += getOrigTarget(i);
            return sum;
        }

        public void setOrigTarget(int targetIndex, double value)
        {
            _origTargets[targetIndex][_model.mInfo.CurrentModelTimeStepIndex, 0] = value;
        }

        /// <summary>Adds a constraint to the problem and returns the constraint index</summary>
        /// <param name="variable">The variable index to apply a weight within the constraint matrix A.</param>
        /// <param name="weight">The weight to apply within the constraint matrix A.</param>
        /// <param name="rhs">The right-hand-side value.</param>
        /// <param name="type">The type of constraint to add.</param>
        public int AddConstraint(int variable, double weight, double rhs, ConstraintType type)
        {
            return AddConstraint(new int[] { variable }, new double[] { weight }, rhs, type);
        }

        /// <summary>Adds a constraint to the problem and returns the constraint index</summary>
        /// <param name="variables">The variable indices to apply weights within the constraint matrix A.</param>
        /// <param name="weights">The weights to apply within the constraint matrix A.</param>
        /// <param name="rhs">The right-hand-side value.</param>
        /// <param name="type">The type of constraint to add.</param>
        public int AddConstraint(int[] variables, double[] weights, double rhs, ConstraintType type)
        {
            if (_A == null)
            {
                _A = new Matrix();
                _b = new Matrix();
                _cTypeList = new List<ConstraintType>();
            }
            int row = _A.NumOfRows;
            _A.AppendRows(1);
            _b.AppendRows(new Matrix(1, 1, rhs));
            for (int i = 0; i < variables.Length; i++)
                _A[row, variables[i]] = weights[i];
            _cTypeList.Add(type);
            return row;
        }

        /// <summary>Adds a variable to the problem with a specified cost and returns the index of the variable.</summary>
        /// <param name="cost">The cost of the variable within the objective function.</param>
        /// <param name="type">The type of variable (negative or not or free)</param>
        public int AddVariable(double cost, VariableType type)
        {
            int col = _c.NumOfRows;
            _A.AppendCols(1);
            _c.AppendRows(1);
            _c[col] = cost;
            _vTypeList.Add(type);
            return col;
        }

        public void Solve(bool linear)
        {
            //for (int i = 0; i < _mu.Length; i++)
            //    _c[i] = _origCosts[i] + _mu[i];
            //_linearSolver = new LinearSolver(_A, _b, _c);
            //_linearSolver.VariableTypes = _vTypeList.ToArray();
            //_linearSolver.ConstraintTypes = _cTypeList.ToArray();
            //_linearSolver.Solve();
            //_decisions = _linearSolver.x;
            //this.SetPowerInjections();
            Solution solution;
            if (linear) 
                solution = _context.Solve(new SimplexDirective());
            else 
                solution = _context.Solve();

            if (solution.Quality != SolverQuality.Optimal)
            {
                _model.FireOnError("Power solver solution quality is: " + solution.Quality.ToString());
                Console.WriteLine("Power solver solution quality is: " + solution.Quality.ToString());
                Console.WriteLine("Solution report: ");
                Console.WriteLine(solution.GetReport());
            }
            List<Decision> decs = new List<Decision>(solution.Decisions);
            _decisions = Array.ConvertAll(decs.ToArray(), dec => dec.ToDouble());
        }

        public Matrix EnergyTargets()
        {
            Matrix m = new Matrix(_targets.Length, 1, 0.0);
            for (int i = 0; i < _targets.Length; i++)
                m[i] = _targets[i].EnergyTarget;
            return m;
        }

        public void SetPowerInjections()
        {
            if (_decisions != null)
            {
                for (int i = 0; i < _targets.Length; i++)
                    _targets[i].EnergyTarget = _decisions[i];
                _model.hydro.Update();
            }
        }

        public Matrix ConstraintValue()
        {
            if (_decisions == null)
                return null;
            Matrix m = new Matrix(_targets.Length, 1, 0.0);
            for (int i = 0; i < _targets.Length; i++)
                m[i] = _targets[i].Energy - _decisions[i];
            return m;
        }

        public bool UpdateMultipliers()
        {
            this.SetPowerInjections();
            return _mu.Update(_model.mInfo.Iteration, this.ConstraintValue());
        }

    }
}
